import os
import numpy as np
import torch
from PIL import Image
from xml.dom.minidom import parse
class week14Dataset(object):
    def __init__(self, root, transforms):
        self.root = root
        self.transforms = transforms
        # load all image files, sorting them to
        # ensure that they are aligned
        self.imgs = list(sorted(os.listdir(os.path.join(root, "week14_tracker_data"))))
        self.anno = list(sorted(os.listdir(os.path.join(root, "week14_tracker_gt/week14-tracker-BWl0oVdzEeUpkHgJTFB/annotation/V001/annotations"))))
   
    def __getitem__(self, idx):
        # load images ad masks
        img_path = os.path.join(self.root, "week14_tracker_data", self.imgs[idx])
        #mask_path = os.path.join(self.root, "GTMasks", self.masks[idx])
        anno_path = os.path.join(self.root,"week14_tracker_gt/week14-tracker-BWl0oVdzEeUpkHgJTFB/annotation/V001/annotations",self.anno[idx])
        print("idx=%s,img_path=%s"%(idx,img_path))
        
        img = Image.open(img_path).convert("RGB")
        
        # 从 anno中读取得到 boxes    
        boxes=[]
        xroot=parse(anno_path).documentElement
        xmin=xroot.getElementsByTagName('xmin')[0].childNodes[0].nodeValue
        ymin=xroot.getElementsByTagName('ymin')[0].childNodes[0].nodeValue
        xmax=xroot.getElementsByTagName('xmax')[0].childNodes[0].nodeValue
        ymax=xroot.getElementsByTagName('ymax')[0].childNodes[0].nodeValue
        boxes.append([int(xmin),int(ymin),int(xmax),int(ymax)])
        num_objs=1
        # convert everything into a torch.Tensor
       
        boxes = torch.as_tensor(boxes, dtype=torch.float32)
        # there is only one class
        labels = torch.ones((num_objs,), dtype=torch.int64)
        #masks = torch.as_tensor(masks, dtype=torch.uint8)

        image_id = torch.tensor([idx])
        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
        # suppose all instances are not crowd
        iscrowd = torch.zeros((num_objs,), dtype=torch.int64)

        target = {}
        target["boxes"] = boxes
        target["labels"] = labels
        #target["masks"] = masks
        target["image_id"] = image_id
        target["area"] = area
        target["iscrowd"] = iscrowd

        if self.transforms is not None:
            img, target = self.transforms(img, target)

        return img, target

    def __len__(self):
        return len(self.imgs)
 
import transforms as T

def get_transform(train):
    transforms = []
    transforms.append(T.ToTensor())
    if train:
        transforms.append(T.RandomHorizontalFlip(0.5))
    return T.Compose(transforms)


def collate_fn(batch):
    return tuple(zip(*batch))



if __name__=="__main__":
    # use our dataset and defined transformations
    dataset = week14Dataset('../week10-dataset/week14_tracker_dataset', get_transform(train=True))
    dataset_test = week14Dataset('../week10-dataset/week14_tracker_dataset', get_transform(train=False))

    # split the dataset in train and test set
    #indices = torch.linspace(0,int(len(dataset))-1,int(len(dataset)),dtype=torch.int).tolist()
    indices = torch.tensor([ii for ii in range(len(dataset))])
    #import pdb
    #pdb.set_trace()
    dataset = torch.utils.data.Subset(dataset, indices[:-50])
    dataset_test = torch.utils.data.Subset(dataset_test, indices[-50:])

    # define training and validation data loaders
    # 用于跟踪的数据，不能打乱，需要按照顺序，所以shuffle=False
    train_loader = torch.utils.data.DataLoader(
        dataset, batch_size=1, shuffle=False, num_workers=1,
        collate_fn=collate_fn)
    for iter, batch in enumerate(train_loader):
        #import pdb
        #pdb.set_trace()
        # batch[0]:image
        # batch[1]:label
        images = batch[0]
        labels = batch[1]
    
        print("batch size :%s"%(len(images)))
        for i,label in enumerate(labels):
            print("label[%s][boxes] :%s"%(i,label['boxes']))
        print("-"*30)